﻿using System;
using System.Linq;

namespace ManagedTFRecordLibrary
{
    internal static class ExampleConverter
    {
        internal static Tensorflow.Example ConvertToTFExample(Example example)
        {
            return new Tensorflow.Example
            {
                Features = new Tensorflow.Features
                {
                    Feature = example.Features.ToDictionary(item => item.Key, item => ConvertToTFFeature(item.Value))
                }
            };
        }

        internal static Example ConvertFromTFExample(Tensorflow.Example example)
        {
            return new Example
            {
                Features = example.Features.Feature.ToDictionary(item => item.Key, item => ConvertFromTFFeature(item.Value))
            };
        }

        internal static Tensorflow.Feature ConvertToTFFeature(Feature feature)
        {
            if (feature is BytesListFeature bytesListFeature)
            {
                return new Tensorflow.Feature
                {
                    BytesList = new Tensorflow.BytesList
                    {
                        Values = bytesListFeature.Values
                    }
                };
            }
            else if (feature is FloatListFeature floatListFeature)
            {
                return new Tensorflow.Feature
                {
                    FloatList = new Tensorflow.FloatList
                    {
                        Values = floatListFeature.Values
                    }
                };
            }
            else if (feature is Int64ListFeature int64ListFeature)
            {
                return new Tensorflow.Feature
                {
                    Int64List = new Tensorflow.Int64List
                    {
                        Values = int64ListFeature.Values
                    }
                };
            }
            throw new NotSupportedException();
        }


        internal static Feature ConvertFromTFFeature(Tensorflow.Feature feature)
        {
            if (feature.ShouldSerializeBytesList())
            {
                return new BytesListFeature(feature.BytesList.Values);
            }
            if (feature.ShouldSerializeFloatList())
            {
                return new FloatListFeature(feature.FloatList.Values);
            }
            if (feature.ShouldSerializeInt64List())
            {
                return new Int64ListFeature(feature.Int64List.Values);
            }
            throw new NotSupportedException();
        }
    }
}
